package org.hepeng.commons.spring.cloud.feign;

import com.netflix.hystrix.strategy.HystrixPlugins;
import com.netflix.hystrix.strategy.concurrency.HystrixConcurrencyStrategy;
import com.netflix.hystrix.strategy.eventnotifier.HystrixEventNotifier;
import com.netflix.hystrix.strategy.executionhook.HystrixCommandExecutionHook;
import com.netflix.hystrix.strategy.metrics.HystrixMetricsPublisher;
import com.netflix.hystrix.strategy.properties.HystrixPropertiesStrategy;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;

import java.util.concurrent.Callable;

/**
 * @author he peng
 */
public class FeignHystrixConcurrencyStrategy extends HystrixConcurrencyStrategy {

    public FeignHystrixConcurrencyStrategy() {
        init();
    }

    private void init() {
        // Keeps references of existing Hystrix plugins.
        HystrixCommandExecutionHook commandExecutionHook =
                HystrixPlugins.getInstance().getCommandExecutionHook();
        HystrixEventNotifier eventNotifier = HystrixPlugins.getInstance().getEventNotifier();
        HystrixMetricsPublisher metricsPublisher = HystrixPlugins.getInstance().getMetricsPublisher();
        HystrixPropertiesStrategy propertiesStrategy =
                HystrixPlugins.getInstance().getPropertiesStrategy();
        HystrixPlugins.reset();
        HystrixPlugins.getInstance().registerConcurrencyStrategy(this);
        HystrixPlugins.getInstance().registerCommandExecutionHook(commandExecutionHook);
        HystrixPlugins.getInstance().registerEventNotifier(eventNotifier);
        HystrixPlugins.getInstance().registerMetricsPublisher(metricsPublisher);
        HystrixPlugins.getInstance().registerPropertiesStrategy(propertiesStrategy);
    }

    @Override
    public <T> Callable<T> wrapCallable(Callable<T> callable) {
        RequestAttributes requestAttributes = RequestContextHolder.getRequestAttributes();
        return new CallableWrapper<>(callable , requestAttributes);
    }

    private static class CallableWrapper<V> implements Callable<V> {

        private final Callable<V> target;
        private final RequestAttributes requestAttributes;

        public CallableWrapper(Callable<V> target, RequestAttributes requestAttributes) {
            this.target = target;
            this.requestAttributes = requestAttributes;
        }

        @Override
        public V call() throws Exception {
            try {
                RequestContextHolder.setRequestAttributes(this.requestAttributes);
                V result = this.target.call();
                RequestContextHolder.resetRequestAttributes();
                return result;
            } finally {
                RequestContextHolder.resetRequestAttributes();
            }
        }
    }
}
